package com.cvai.service.impl;

import com.cvai.aiservice.FileChatAssistant;
import com.cvai.aiservice.KnowledgeChatAssistant;
import com.cvai.context.BaseContext;
import com.cvai.domain.dto.KnowledgeDto;
import com.cvai.domain.po.Chat;
import com.cvai.domain.po.File;
import com.cvai.domain.po.Knowledge;
import com.cvai.domain.po.Memory;
import com.cvai.exception.CvAiException;
import com.cvai.mapper.KnowledgeMapper;
import com.cvai.service.IFileService;
import com.cvai.service.IKnowledgeService;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByCharacterSplitter;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.transformer.CompressingQueryTransformer;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicReference;

/**
 * <p>
 *  服务实现类
 * </p>
 *
 * @author cv
 * @since 2025-03-14
 */
@Service
@RequiredArgsConstructor
@Slf4j
public class KnowledgeServiceImpl extends ServiceImpl<KnowledgeMapper, Knowledge> implements IKnowledgeService {
    private final IFileService fileService;
    private final QdrantClient qdrantClient;
    private final EmbeddingModel embeddingModel;
    private final ChatLanguageModel chatLanguageModel;
    private final StreamingChatLanguageModel streamingChatLanguageModel;
    @Override
    @Transactional(rollbackFor = RuntimeException.class)
    public void add(KnowledgeDto knowledgeDto) {
        Knowledge knowledge = new Knowledge()
                .setName(knowledgeDto.getName())
                .setUserId(BaseContext.getCurrentUser().getId());
        if(getOne(lambdaQuery().getWrapper().eq(Knowledge::getName,knowledge.getName()))!=null)
            throw new CvAiException("知识库已存在");
        save(knowledge);
        DocumentParser parser = new ApacheTikaDocumentParser();

        List<Path> paths = new ArrayList<>();
        knowledgeDto.getFileIds().forEach(fileId -> {
            File file = fileService.getById(fileId);
            if(file==null)
                throw new CvAiException("文件不存在");
            if(file.getIsPicture()==1)
                throw new CvAiException("图片文件不能作为知识库");

            Path path = Paths.get(file.getFilePath());
            paths.add(path);
        });

        var vectorParams = Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(1024) //这个是根据llm的维度决定的
                .build();
        try {
            qdrantClient.createCollectionAsync(knowledge.getName(), vectorParams).get();//创建集合
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } catch (ExecutionException e) {
            throw new RuntimeException(e);
        }

        QdrantEmbeddingStore store = QdrantEmbeddingStore.builder()
                .client(qdrantClient)
                .collectionName(knowledge.getName())
                .build();
        DocumentByCharacterSplitter splitter = new DocumentByCharacterSplitter(100, 20);
        paths.forEach(path -> {
            Document document = FileSystemDocumentLoader.loadDocument(path, parser);
            splitter.split(document).forEach(textSegment -> store.add(embeddingModel.embed(textSegment).content(), textSegment));
        });
    }

    @Override
    @Transactional
    public void delById(Long id) {
        qdrantClient.deleteCollectionAsync(getById(id).getName());
        removeById(id);
    }

    @Override
    public Flux<String> chat(Long id, String question) {
        Knowledge knowledge = getById(id);
        if(knowledge ==null)
            throw new RuntimeException("知识库不存在");
        QdrantEmbeddingStore embeddingStore = QdrantEmbeddingStore.builder()
                .client(qdrantClient)
                .collectionName(knowledge.getName())
                .build();

        EmbeddingStoreContentRetriever retriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(10)
                .minScore(0.5)
                .build();

        CompressingQueryTransformer queryTransformer = new CompressingQueryTransformer(chatLanguageModel);

        KnowledgeChatAssistant knowledgeChatAssistant = AiServices.builder(KnowledgeChatAssistant.class)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(5))
                .chatLanguageModel(chatLanguageModel)
                .streamingChatLanguageModel(streamingChatLanguageModel)
                .retrievalAugmentor(DefaultRetrievalAugmentor.builder()
                        .contentRetriever(retriever)
//                        .queryTransformer(queryTransformer)
                        .contentInjector(new DefaultContentInjector())
                        .build()).build();
        AtomicReference<StringBuilder> fullContent = new AtomicReference<>(new StringBuilder());

        return knowledgeChatAssistant.chat(UUID.randomUUID().toString(), question)
                // 1. 处理每个数据块（发送到前端 + 本地收集）
                .doOnNext(s -> {
                    fullContent.get().append(s);
                    log.info("收到流式块: {}", s);
                })
                //2. 格式化为sse 格式
                // .map(s -> "data: " + s + "\n\n")
                //3. 流结束时异步保存数据
                .concatWith(Mono.just("end"))
                .doOnComplete(() -> {
                    String ans = fullContent.get().toString();
                    log.info("收到完整的流式块: {}", ans);
                    if (!ans.isEmpty()) {
                        Mono.fromRunnable(() -> {

                                }).subscribeOn(Schedulers.boundedElastic()) // 切换到阻塞线程池
                                .subscribe(
                                        null,
                                        error -> log.error("存储失败", error)
                                );
                    }
                });
    }

}
